package com.mingserve.common.core.utils;

import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.transform.DftNormalization;
import org.apache.commons.math3.transform.FastFourierTransformer;
import org.apache.commons.math3.transform.TransformType;

import java.util.Arrays;

/**
 * 信号处理相关工具类
 * 使用Apache Commons Math库 (org.apache.commons:commons-math3)FastFourierTransformer进行FFT计算
 *
 * @author lvwshuai
 * @date 2023-04-04
 */
public class SignalUtils2
{
    /**
     * 按采样频率生成采集时间间隔
     *
     * @param n          数据长度
     * @param sampleRate 采样频率
     * @return double[] 时间间隔保留6位小数
     */
    public static double[] intervals(int n, double sampleRate)
    {
        double interval = 1 / sampleRate; // 每个数据点的间隔
        double[] intervals = new double[n];
        for (int i = 0; i < n; i++) {
            intervals[i] = MathUtils.round(i * interval, 6);
        }
        return intervals;
    }

    /**
     * 生成一个长度为n/2的数组，表示在频率范围[0, fs/2]内每个频率上的幅值。
     *
     * @param n          数据长度
     * @param sampleRate 采样频率
     * @return double[] 频率保留2位小数
     */
    public static double[] freq(int n, double sampleRate)
    {
        double[] freq = new double[n / 2 + 1];
        for (int i = 0; i <= n / 2; i++) {
            freq[i] = MathUtils.round(i * sampleRate / n, 2);
        }
        return freq;
    }

    /**
     * 去直流分量 (DC滤波)
     *
     * @param signal 原始信号
     * @return 原始信号减去平均值
     */
    public static double[] removeDC(double[] signal)
    {
        if (signal == null) return null;
        double mean = Arrays.stream(signal).average().orElse(0);
        return Arrays.stream(signal).map(s -> s - mean).toArray();
    }

    /**
     * 获取振动信号的FFT
     * 使用Apache Commons Math库 (org.apache.commons:commons-math3)FastFourierTransformer进行FFT计算
     *
     * @param vib 振动信号 2的幂指数值
     * @param Fs  采样频率
     * @return double[][]{vib_fft保留5位小数, f保留2位小数}
     */
    public static double[][] fft(double[] vib, double Fs)
    {
        if(vib == null || vib.length == 0) return null;

        vib = removeDC(vib);

        // 使用FastFourierTransformer进行FFT
        FastFourierTransformer fft = new FastFourierTransformer(DftNormalization.STANDARD);
        Complex[] vibFftComplex = fft.transform(vib, TransformType.FORWARD);

        int N = vib.length;

        // 计算频谱幅值和频率
        double[] vib_fft = new double[N / 2 + 1];
        for (int i = 0; i <= N / 2; i++) {
            vib_fft[i] = MathUtils.round(2.0 * vibFftComplex[i].abs() / N,5);
        }

        double[] f = new double[N / 2 + 1];
        for (int i = 0; i <= N / 2; i++) {
            // Fs / 2 表示奈奎斯特频率（Nyquist Frequency
            // (double) i / (N / 2) 将索引 i 映射到 [0, 1] 区间
            // 结果等价于 i * Fs / N
            // f[i] = MathUtils.round(Fs / 2 * ((double) i / (N / 2)),2);
            f[i] = MathUtils.round(i * Fs / N,2);
        }

        return new double[][]{vib_fft, f};
    }

    /**
     * 计算振动信息的包络信号
     * 使用Apache Commons Math库 (org.apache.commons:commons-math3)FastFourierTransformer进行FFT计算
     *
     * @param vib 振动信号
     * @return double[] 包络信号
     */
    public static double[] env(double[] vib)
    {
        int N = vib.length;
        // 转换为Complex数组，并构建复数数组
        Complex[] vibComplex = new Complex[N];
        for (int i = 0; i < N; i++) {
            vibComplex[i] = new Complex(vib[i], 0);
        }

        // 使用FastFourierTransformer进行FFT
        FastFourierTransformer fft = new FastFourierTransformer(DftNormalization.STANDARD);
        Complex[] transformed = fft.transform(vibComplex, TransformType.FORWARD);

        // 构建Hilbert变换，保留正频率分量并置零负频率分量
        for (int i = 1; i < N / 2; i++) {
            transformed[i] = transformed[i].multiply(2);
        }
        for (int i = N / 2 + 1; i < N; i++) {
            transformed[i] = Complex.ZERO;
        }

        // 进行逆FFT
        Complex[] inverseTransformed = fft.transform(transformed, TransformType.INVERSE);

        // 计算包络信号
        double[] vibEnvelope = new double[N];
        for (int i = 0; i < N; i++) {
            vibEnvelope[i] = inverseTransformed[i].abs();
        }

        return vibEnvelope;
    }
}
